program molsem
  !
  ! ML/OS-EM Loop Handler
  !
  !
  ! LICENSE AGREEMENT:
  !
  ! You agree that we don't agree to anything.
  !
  !
  ! DISCLAIMER OF WARRANTIES:
  !
  ! This software is provided to you "AS IS," and the authors and
  ! contributors disclaim any warranty or liability obligations to you
  ! of any kind, implied or express, including but not limited to the
  ! warranties of merchantability, fitness for a particular purpose,
  ! and non-infringement.
  ! 
#define green(i) (iand(mana,i)/=0)
  use omp_lib 
  use suty
  use kuty
  use vuty
  use duty
  use ground_control
  use PLANAR_PEM
  use projectors
  implicit none
  !
  real,parameter::eps=epsilon(1.0)
  integer opt
  real*4,allocatable,target::x_js(:,:,:,:)
  real*4,pointer::x_j(:,:,:)
  type(r4v3),allocatable::C_js(:)
  real*4,pointer::C_j(:,:,:)
  real*4,allocatable,target::D_js(:,:,:,:)
  real*4,pointer::D_j(:,:,:)
  real*4,allocatable,target::p_is(:,:,:,:,:)
  real*4,pointer::p_i(:,:,:,:)
  real*4,allocatable,target::q_is(:,:,:,:,:)
  real*4,pointer::q_i(:,:,:,:)
  real*4,allocatable,target::y_is(:,:,:,:,:)
  real*4,pointer::y_i(:,:,:,:)
  real*4,allocatable::u_i(:,:,:,:) ! normalize
  real*4,allocatable::C_i(:,:,:,:) ! 
  real,parameter::mu_water=9.607e-3 ! 1/mm

  !
  integer ite,is,isub
  real logL,logL_old,maxlogL
  integer istat
  integer::submap(4)=0
  real::yi_scale=1 ! to make avegt0 p_i approx unity
  real::xj_scale=1 ! to make avegt0 x_j approx unity
  integer ni
  real t0,t1
  integer rts(8,2)
  integer ires
  type(cmd_info) ci
  type(image_info) ifo
  real th
  real calib,retr

  !
  ite=-1
  is=0
  calib=1
  call set_rstat(-10)
  call me_first_mana(ci)
  opt=mana
  ci%pem%option=opt
  if(ci%num_subsets>1) submap=[1,2,3,4]
  ires=2
  call echo("Misson code registered: "//trim(ci%mc))

  !
  call set_rstat(-9)
  call get_einfo(ifo)
  call gc2mana("initnializing modlue `planar_pem'")
  call setenv_ppem(ci%pem,ires)
  call printenv_ppem()
  allocate(y_is(mx,my,mx,my,2),p_is(mx,my,mx,my,2),q_is(mx,my,mx,my,2))
  p_i=>p_is(:,:,:,:,1) 
  q_i=>q_is(:,:,:,:,1) 
  y_i=>y_is(:,:,:,:,1) 
  if(ci%ofmtid/=OFMT_RAW) call set_idim(ifo,MAT_IMG)

  !
  call set_rstat(-8)
  call gc2mana("loading input data")
  if(.not.green(OPT_SYSMAT_ONLY)) then
     if(ci%nfiles(FT_LIST)>0) then
        call read_list(ci%narg,y_i,indx=ci%argindxfile)
        if(green(OPT_LMD_FLIPB1)) call flipHistB1(y_i)
        if(green(OPT_LMD_FLIPX1)) call flipHistX1(y_i)
        if(green(OPT_LMD_FLIPX2)) call flipHistX2(y_i)
        if(green(OPT_LMD_FLIPY1)) call flipHistY1(y_i)
     end if
     if(ci%nfiles(FT_HIST)>0) call read_hist(ci%narg,y_i,indx=ci%argindxfile)
     yi_scale=sum(y_i)
     if(yi_scale==0) then
        call echo("*** Warning: histogram is empty.")
     else
        xj_scale=real(mx*my_true)**2/yi_scale
        yi_scale=real(nv_true)/yi_scale
     end if
     call echo_v(y_i,"y")
  end if

  !
  call set_rstat(-7)
  call gc2mana("evaluating system matrix")
  call set_sysmat(yi_scale)
  call resort_sysmat

  !
  call set_rstat(-6)
  call gc2mana("initializing module `projectors'")
  call alloc_Cjs_aligned
  call alloc_xjs_aligned
  !
  x_j=>x_js(:,:,1:nz,1)
  D_j=>D_js(:,:,1:nz,1)
  !
  call init_proj(ires,maxval(sysmat%n),int(omp_get_max_threads()),&
       ci%num_subsets,option=mana)

  !
  call set_rstat(-5)
  call gc2mana("generating 2d image")
  call bproj_2d(y_i,x_j,ci%pem%detector_separation) ! <<<
  call echo_v(x_j(1:nx,1:ny_true,1:nz),"x_j")
  isub=-1
  if(green(OPT_WRITE_IMAG)) call out_img(x_j,"2d")
 
  !
  call set_rstat(-4)
  call gc2mana("resorting histogram")
  call tros_hist(y_i)

  !
  call set_rstat(-3)
  call gc2mana("evaluating Cj")
  ! if Cj be loaded, OPT_NORMALIZE is ignored
  if(green(OPT_LOAD_Cj)) then
     call load_cj(ci%fcj,C_j)     
  else
     allocate(u_i(mx,my,mx,my),C_i(mx,my,mx,my))
     if(green(OPT_NORMALIZE)) then
        call read_hist(1,u_i,file=ci%fnorm)
     else
        u_i=1.0_4
     end if
     if(mg/=0) call set_hist_mask(u_i)
     call tros_hist(u_i)
     call echo_v(u_i,"Raw u_i")
     if(green(OPT_LOAD_CI)) then
        call load_ci(ci%fci,C_i)
     else
        x_j=1.0_4
        call fproj(x_j,C_i)
        if(green(OPT_DUMP_CI)) call write_v(trim(ci%foutb)//"-Ci",-1,-1,loc(C_i),size(C_i),istat,ci%ofmtid)
        if(green(OPT_SAVE_CI)) call save_Ci(ci%fci,C_i)
     end if
     call echo_v(C_i,"C_i")
     where(C_i>0)
        u_i=u_i/C_i
     elsewhere
        u_i=0.0_4
     end where
     deallocate(C_i)
     if(ci%pem%thi>0) then
        th=maxval(u_i)*ci%pem%thi
        call echo("Ci threashold: "//trim(f2a(th,-1)))
        where(u_i<th) u_i=0.0_4
     end if     
     where(u_i==0) y_i=0.0_4
     if(green(OPT_DUMP_UI)) call write_v(trim(ci%foutb)//"-Ui",-1,-1,loc(u_i),size(u_i),istat,ci%ofmtid)
     call echo_v(u_i,"Adj u_i")
     do is=1,ci%num_subsets
        isub=submap(is)
        call init_subset(isub)
        C_j=>C_js(is)%a
        call bproj(u_i,C_j)
        if(green(OPT_SAVE_CJ)) call save_Cj(ci%fcj,C_j) ! incomplete
     end do
  end if
 
  if(ci%pem%thj>0) then
     th=maxval(C_j)*ci%pem%thj
     call echo("Cj threashold: "//trim(f2a(th,-1)))
     where(C_j<th) C_j=0.0_4
  end if
  C_j=C_j/real(xj_scale,4)
  call echo("xj scaling factor: "//trim(f2a(xj_scale,-1)))
  call echo_v(C_j,"C_j")
  
  !
  if(green(OPT_SYSMAT_ONLY)) stop "done."

  !
  call set_rstat(-2)
  call gc2mana("initializing image")
  if(green(OPT_USE_LOGL)) then
     p_i=y_i
     logL=0
     maxlogL=get_LogL()
     call echo("Max Log L = "//trim(f2a(maxLogL,-1)))
     p_i=0.0_4
  end if
  call init_xj ! set initial guess

  !
  calib=ci%pem%cf
  if(green(OPT_DECAY_CORRECTION)) then
     retr=getDecayCorrectionFactor()
     if(retr<=0) call echo_red("Failed to obtain decay correction factor")
     call echo("Decay correction factor: "//trim(f2a(retr,4)))
     calib=retr
  end if
  if(green(OPT_CALIBRATION)) then
     retr=getCalibrationFactor()
     if(retr<=0) call echo_red("Failed to obtain calibration factor")
     call echo("Dose/Weight calibration factor: "//trim(f2a(retr,4)))
     calib=calib*retr
  end if
  call echo("Calibration factor: "//trim(f2a(calib,4)))
  
  call set_rstat(-1)  
  call gc2mana("MLEM about to begin")
  
  !
  ite=-1
  if(green(OPT_WRITE_CJ)) call out_img(C_j,"Cj",.true.)
  ite=0
  call out_imgs

  !
  call set_rstat(RSTAT_RUN)
  call echo("MLEM Begin.")

  !
  ni=0
  call cpu_time(t0)
  call date_and_time(values=rts(:,1))
  
  mlem: do ite=1,ci%num_iteration_max
     call set_rstat(ITE)
     call gc2mana
     subset: do is=1,ci%num_subsets
        isub=submap(is)
        C_j=>C_js(is)%a 
        call init_subset(isub)
        ni=ni+1

        !
        call fproj(x_j,p_i)
        call divps(loc(q_i),loc(y_i),mxy2,loc(p_i))
        call print_stat_proj
        call out_projs

        !
        call bproj(q_i,D_j)
        call muldivps(loc(x_j),loc(D_j),nv,loc(C_j))

        call print_stat_imag
        call out_imgs

     end do subset
     !
     if(green(OPT_USE_LOGL)) then
        if(isConverge()) exit mlem
     end if
     !     
  end do mlem

  !
  call set_rstat(0)

  !
  call date_and_time(values=rts(:,2))
  call cpu_time(t1)

  !  
  call echo("Elapsed time: "//trim(f2a((t1-t0)/ni,4))//"u "&
       //trim(f2a(etime(rts(:,1),rts(:,2))/ni,4))//"s")

  !
  call set_rstat(RSTAT_DONE)
  
contains
  !
  subroutine init_xj
    real sCj,xj0
    real*4,allocatable,target::Cj_mask(:,:,:)
    integer is
    if(green(OPT_UNIFORM_INIT)) then
       call echo("Using uniform image as an initial guess.")
       x_j=1.0_4
       return
    end if
    allocate(Cj_mask(nx,ny,nz))
    sCj=0
    do is=1,ci%num_subsets
       C_j=>C_js(is)%a
       sCj=sCj+sum(C_j)
       Cj_mask=Cj_mask+C_j
    end do
    where(Cj_mask==0) x_j=0.0_4
    xj0=sum(y_i)/sCj
    x_j=x_j*real(xj0,4)/sum(x_j)
    if(green(OPT_VERBOSE)) &
         call echo("Initial guess: "//trim(f2a(xj0,-1)))
    deallocate(Cj_mask)
    if(green(OPT_USE_FILTER)) then
       call conv_img(x_j,.true.)
    end if
  end subroutine init_xj
  !
  subroutine alloc_xjs_aligned
    integer dif
    dif=nz-4*(nz/4)
    if(dif/=0) dif=4-dif
    allocate(D_js(nx,ny,nz+dif,2),x_js(nx,ny,nz+dif,2))
  end subroutine alloc_xjs_aligned
  !
  subroutine alloc_Cjs_aligned
    if(allocated(C_js)) then
       do is=1,size(C_js)
          deallocate(C_js(is)%a)
       end do
       deallocate(C_js)
    end if
    allocate(C_js(ci%num_subsets))
    C_j=>C_js(1)%a ! <<<
    do is=1,ci%num_subsets
       allocate(C_js(is)%a(nx,ny,nz))
    end do
  end subroutine alloc_Cjs_aligned
  !
  real function get_LogL()
    real*4 v(mx,my,mx,my)
    v=p_i
    where(v>0)
       v=-(v)+y_i*log(v)
    elsewhere
       v=0.0_4
    end where
    get_logL=sum(v(1:mx,1:my_true,1:mx,1:my_true))
  end function get_LogL
  !
  logical function isConverge()
    isConverge=.false.
    if(green(OPT_USE_OSEM)) call fproj(x_j,p_i)
    logL_old=logL
    logL=get_LogL()
    if(green(OPT_VERBOSE)) then
       call echo("ite, logL = "//trim(i2a(ite))//", "//trim(f2a(logL,-1))//" ("&
            //trim(f2a(logL/maxLogL*100,1))//" %)")
    end if
    if(abs(logL-logL_old)/abs(logL)<=eps) isConverge=.true.
  end function isConverge
  !
  subroutine print_stat_proj
    call echo_v(p_i(1:mx,1:my_true,1:mx,1:my_true),"p")
    call echo_v(q_i(1:mx,1:my_true,1:mx,1:my_true),"q")
  end subroutine print_stat_proj
  !
  subroutine print_stat_imag
    call echo_v(d_j(1:nx,1:ny_true,1:nz),"D_j")
    call echo_v(x_j(1:nx,1:ny_true,1:nz),"x_j")
  end subroutine print_stat_imag 
  !
  subroutine out_imgs
    if(green(OPT_WRITE_UPDT).and.ite>0) call out_img(D_j,"Dj",.true.)
    if((green(OPT_WRITE_IMAG).and.ite>=0)&
         .or.ite==ci%num_iteration_max) call out_img(x_j,"")
  end subroutine out_imgs
  !
  subroutine out_img(ov,ap,onfile)
    real*4,intent(in)::ov(:,:,:)
    character*(*),intent(in)::ap
    logical,intent(in),optional::onfile
    real*4 img(nx,ny_true,nz)
    integer jx,jy,jz
    integer pimg,simg
    integer i1
    character(MAXPATHLEN) fout
    if(ny/=ny_true) then
       forall(jx=1:nx,jy=1:ny_true,jz=1:nz)
          img(jx,jy,jz)=real(ov(jx,jy,jz)*calib,4)
       end forall
       simg=nx*ny_true*nz
       pimg=loc(img)
    else
       simg=nv
       if(calib==1) then
          pimg=loc(ov)
       else
          img=real(ov*calib,4)
          pimg=loc(pimg)
       end if
    end if
    if(ite==ci%num_iteration_max) then
       i1=-1
    else
       i1=ite
    end if
    fout=trim(ci%foutb)
    if(ap/="") fout=trim(fout)//"-"//ap
    call write_v(fout,i1,isub,pimg,simg,istat,ci%ofmtid,ifo)
    if(shmstats(SHMID_OUT)%p>0.and..not.present(onfile)) call write_v_shm(pimg,ifo)
  end subroutine out_img

  !
  subroutine out_projs
    if(green(OPT_DUMP_PROJ)) then
       call write_v(trim(ci%foutb)//"-pi",ite,isub-1,loc(p_i),size(p_i),istat,ci%ofmtid)
       call write_v(trim(ci%foutb)//"-qi",ite,isub-1,loc(q_i),size(q_i),istat,ci%ofmtid)
    end if
  end subroutine out_projs

  !
  subroutine echo(m)
    character*(*),intent(in)::m
    if(.not.green(OPT_VERBOSE)) return
    write(*,10) trim(m)
    call flush(6)
10  format(x,"(ml) ",a)
  end subroutine echo

  real function getDecayCorrectionFactor()
    character*128 str
    real t0,t1,dt,tau,x
    getDecayCorrectionFactor=-1
    call getenv("RP3SINT",str)
    if(len_trim(str)/=5) return
    t0=(a2f(str(1:2))*60+a2f(str(4:5)))*60

    call getenv("RP3SMINUTE",str)
    if(len_trim(str)/=2) return
    t1=a2f(str)*60

    call getenv("RP3SHOUR",str)
    if(len_trim(str)/=2) return
    t1=t1+a2f(str)*60*60-t0

    call getenv("RP3SSCAND",str) ! in sec
    if(str=="") return
    dt=a2f(str)

    call getenv("RP3HALFLIFE",str) ! in min
    if(str=="") return
    tau=a2f(str)*60/log(2.0_dp)

    x=dt/tau
    getDecayCorrectionFactor=x*exp(t1/tau)/(1.0_dp-exp(-x))

  end function getDecayCorrectionFactor

  real function getCalibrationFactor()
    character*128 str
    real w,d,p
    getCalibrationFactor=-1
    
    call getenv("RP3PWEIGHT",str) ! kg
    if(str=="") return
    w=a2f(str)

    call getenv("RP3SDOSE",str) ! Bq
    if(str=="") return
    d=a2f(str)/3.7e10_dp*1.0e3_dp ! mCi
    
    p=1
    call getenv("RP3PFRAC",str)
    if(str/="") then
       p=a2f(str)
    else
       call echo("*** getCalibrationFactor: RP3PFRAC not set, assuming Positron fraction = 1.0")
    end if

    getCalibrationFactor=p/(d/w)

  end function getCalibrationFactor
  
  !
  subroutine echo_red(mess,code)
    character*(*),intent(in)::mess
    integer,intent(in),optional::code
    integer ec
    character*(*),parameter::swan_song="*** emergency escape from main."
    write(0,*) trim(adjustl(mess))
    call set_rstat(RSTAT_ERROR)
    ec=-1
    if(present(code)) ec=code
    call exit(ec)
  end subroutine echo_red

end program molsem
